from decoder import *
import copy

class MemoryBank(nn.Module):
    def __init__(self, test=False, task_nums = 4):
        super(MemoryBank, self).__init__()
        if test:
            self.memory = nn.ModuleList([
                Decoder() for _ in range(task_nums)
            ])
        else:
            self.memory = nn.ModuleList([
                Decoder() 
            ])

    def forward(self, x, t_id):
        # 复制权重并进行训练 
        if t_id >= len(self.memory):
            self.memory.append(copy.deepcopy(self.memory[-1]))

        output = self.memory[t_id](x)
        return output


if __name__ == '__main__':
    mm = MemoryBank(4)
    print(mm)

